#!/usr/bin/env python

import os
import sys
import argparse
import flowtool
import http_tls
import http_tls_40
import sites_plugin
import cPickle as pickle
import setup_environment
import metadata_attributes
import add_plugin


parser = argparse.ArgumentParser(
    """
    This program includes all kinds of QA testing scripts.
    NOTE: Please edit setup_environment.py file for flowtool path and config file before using flowtool option.
    NOTE: Please use setup option before use other options.


    options:
        setup is to setup environment variables.
        flowtool is running flowtool against a single pcap file or a pcap folder.
        web is checking flowtool log files for miss classified http tls tcp udp traffic.
        sites is used for checking format of guids and rules or making webpatterns

    Sample usage:
    ./python main.py setup /root/builds/
    ./python main.py setup -f /root/flowtool/config.txt /root/builds/

    ./python main.py flowtool -c=/root/build/navl_4.0/tool/capreader/ /root/traces/netflix/
    ./python main.py flowtool -c=/root/build/navl_4.0/tool/capreader/ /root/traces/netflix/netlifx.pcap

    ./python main.py web -v 4.0 -l /root/traces/netflix/ /root/results/missed_tls_http_tcp.txt
    """)
subparsers = parser.add_subparsers(dest = 'function'
                                    , help='commands')

"""Setup"""
setup_parser = subparsers.add_parser('setup', help = 'Add plugins to a NAVL build')
setup_parser.add_argument('-f'
                        , '--flowtool'
                        , help = 'Config flowtool path and configuration file path'
                        , metavar = 'config')
setup_parser.add_argument('build'
                        , help = 'Copy QClassify files. Make qclassify and capreader'
                        , metavar = 'build_path')

"""Plugin"""
plugin_parser = subparsers.add_parser('plugin', help = 'Add plugins to a NAVL build.')
plugin_parser.add_argument('-p'
                        , '--plugin'
                        , help = 'Plugin directory path'
                        , metavar = 'plugins_dir'
                        , required = True)
plugin_parser.add_argument('build'
                        , help = 'Navl build path'
                        , metavar = 'navl_build')

"""Flowtool"""
flowtool_parser = subparsers.add_parser('flowtool', help='Run flowtool against pcap file or folder')
flowtool_parser.add_argument('-c'
                            , '--capreader'
                            , help = 'Specify capreader path'
                            , metavar = 'capreader'
                            , required = True)
flowtool_parser.add_argument('-v'
                            , '--version'
                            , help = 'Specify capreader version number, example: 3.0 or 4.0'
                            , metavar = 'version')
flowtool_parser.add_argument('-o'
                            , '--output'
                            , help = 'Specify output folder path'
                            , metavar = 'output')
flowtool_parser.add_argument('pcap'
                            , help = 'Specify pcap file path'
                            , metavar = 'pcap')

"""HTTP, TLS, TCP"""
web_parser = subparsers.add_parser('web', help='Filter unclassified HTTP, HLS and TCP.')
web_parser.add_argument('-v'
                        , '--version'
                        , help = 'Specify capreader version number, example: 3.0 or 4.0'
                        , metavar = 'version'
                        , default = '4.0'
                        , required = True)
web_parser.add_argument('-l'
                        , '--log'
                        , help = 'Specify log file path'
                        , metavar = 'log'
                        , required = True)
web_parser.add_argument('output'
                        , help = 'Specify output file name'
                        , metavar = 'output')

"""Metadata"""
metadata_parser = subparsers.add_parser('metadata'
                                        , help='Extract specific metadata attributes from NAVL 4.0 flowtool log files')
metadata_parser.add_argument('-l'
                        , '--log'
                        , help = 'Specify log file folder path'
                        , metavar = 'log'
                        , required = True)
metadata_parser.add_argument('output'
                        , help = 'Specify output file name'
                        , metavar = 'output')


"""Sites Plugin Related"""
sites_parser = subparsers.add_parser('sites', help = 'Sites related functions.')
sites_parser.add_argument('-g'
                        , '--guid'
                        , help = 'Specify file contains the guid you want check'
                        , metavar = 'guid')
sites_parser.add_argument('-r'
                        , '--rule'
                        , help = 'Specify file contains the rule you want check'
                        , metavar = 'rule')
sites_parser.add_argument('-m'
                        , '--make_sites'
                        , help = 'Make web patterns from guid and rule file'
                        , metavar = ("guid", "rule")
                        , nargs = 2)
sites_parser.add_argument('-M'
                        , '--make_referer'
                        , help = 'Specify '
                        , metavar = ("guid", "rule")
                        , nargs = 2)

args = parser.parse_args()


def check_version(capreader = str):
    if str(capreader).find('4.0') >= 0:
        return '4.0'
    elif str(capreader).find('3.0') >= 0:
        return '3.0'
    return None


def main():
    if args.function == 'setup':
        if args.config:
            if os.path.isfile(args.config):
                conf = setup_environment.Config(args.config)
                output = open("conf.p", "wb")
                pickle.dump(conf, output)
                output.close()
            else:
                print 'ERROR in main(): Invalid configuration file!'
        elif args.build:
            setup_environment.Environment(args.build)
        else:
            print 'ERROR in main(): Invalid build path!'

    elif args.function == 'plugin':
        if os.path.isdir(args.plugin) and os.path.isfile(args.build):
            plugins = add_plugin.Plugins(args.build, args.plugin)
            plugins.extract_build()
            plugins.extract_plugins()
        else:
            print 'ERROR in main(): Invalid arguments!'

    elif args.function == 'flowtool':
        ft = flowtool.Flowtool(args.capreader)
        if args.version == '4.0' or check_version(args.capreader) == '4.0':
            print 'Flowtool 4.0'
            if os.path.isfile(args.pcap):
                ft.run_flowtool_40(args.pcap, args.output)
            elif os.path.isdir(args.pcap):
                ft.run_flowtool_folder_40(args.pcap, args.output)
            else:
                print 'ERROR in main(): Invalid file or folder!'
        elif args.version == '3.0' or check_version(args.capreader) == '3.0':
            print 'Flowtool 3.0'
            if os.path.isfile(args.pcap):
                ft.run_flowtool(args.pcap)
            elif os.path.isdir(args.pcap):
                ft.run_flowtool_folder(args.pcap)
            else:
                print 'ERROR in main(): Invalid file or folder!'
        else:
            print 'ERROR in main(): Invalid capreader version number!'

    elif args.function == 'web':
        if os.path.isdir(args.log):
            if args.version == '4.0':
                print 'Web 4.0'
                if args.output is None:
                    print 'ERROR in main(): No output path specified! Use -o to specify output path.'
                    exit(1)
                ht = http_tls_40.HttpTls()
                ht.missed_http_tls_tcp_folder(args.log, args.output)
            elif args.version == '3.0':
                print 'Web 3.0'
                ht = http_tls.HttpTls()
                ht.missed_http_tls_tcp_folder(args.log, args.output)
            else:
                print 'ERROR in main(): Invalid capreader version number!'
        else:
            print 'ERROR in main(): Invalid log folder path!'

    elif args.function == 'metadata':
        if os.path.isdir(args.log) and os.path.isfile(args.output):
            metadata = metadata_attributes.MetadataAttributes()
            metadata.metadata_folder(args.log, args.output)
        else:
            print 'ERROR in main(): Invalid argument in metadata option!'

    elif args.function == 'sites':
        sites = sites_plugin.SitesPlugin()
        if args.guid:
            sites.check_guid_format(args.guid)
        elif args.rule:
            sites.check_rule_format(args.rule)
        elif args.make_sites:
            sites.make_webpattern_sites(args.make_sites[0], args.make_sites[1])
        elif args.make_referer:
            sites.make_webpattern_referer_sites(args.make_referer[0], args.make_referer[1])
        else:
            print 'ERROR in main(): Invalid argument in sites option!'


if __name__ == "__main__":
    main()
